# Copyright (C) 2009 Frank Ronny Larsen.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from scapy import *
import threading
import time
import gobject


from httpchat import HTTP_Chat


class SnoopThread(threading.Thread):
    def __init__(self, liststore, filename=None, main=None):
        super(SnoopThread, self).__init__()
        self.liststore = liststore
        self.quitting = False
        self.chats = dict()
        self.filename = filename
        self.main = main
        self.setName("SnoopThread")

    def quit(self):
        # TODO: When sniff() gets a stop method to watch for put that
        # here.
        self.quitting = True

    def pkt_callback(self, pkt):
        src = "%s:%s" % (pkt[IP].src, pkt[TCP].sport)
        dst = "%s:%s" % (pkt[IP].dst, pkt[TCP].dport)
        # Create hash of sorted src-dst
        ret = "-".join(sorted([src, dst]))
        if ret in self.chats:
            # Append to existing chat
            chat = self.chats[ret]
            ret = chat.append(pkt, src, dst)
        else:
            # Create new.
            chat = HTTP_Chat(pkt, store=self.liststore)
            self.chats[ret] = chat

    def myfilter(self, pkt):
        # Only practical way of stopping the snooper...
        if self.quitting:
            raise KeyboardInterrupt

        if TCP not in pkt:
            return False

        # Check if connection is in the dict.
        src = "%s:%s" % (pkt[IP].src, pkt[TCP].sport)
        dst = "%s:%s" % (pkt[IP].dst, pkt[TCP].dport)
        # Create hash of sorted src-dst
        index = "-".join(sorted([src, dst]))
        if index in self.chats:
            return True

        # If we don't get a raw-frame, no use continuing.
        if Raw not in pkt:
            return False

        # Check for HTTP.  This will trap if this is a new HTTP
        # connection.
        if "HTTP" in pkt[Raw].load:
            return True

        return False

    def run(self):
        print "Thread starting"
        ret = True

        params = { "prn": self.pkt_callback,
                   "store": 0,
                   "lfilter": self.myfilter }

        if self.filename:
            params["offline"] = self.filename

        try:
            sniff(**params)
        except Exception, ex:
            print "Sniffer failed!", ex
            gobject.idle_add(self.main.snooperQuit, False, ex)
        else:
            gobject.idle_add(self.main.snooperQuit, True)
        
        print "Thread exiting"
